package ctxdb

import (
	"context"
	"github.com/gocraft/dbr/v2"
	"github.com/pkg/errors"
)

// CtxDB extend dbr.Connection with NewContextWithTx, Tx, Commit, CommitOrAbort, Rollback, RollbackUnlessCommitted
type CtxDB struct {
	*dbr.Connection

	Key interface{}
}

type CtxKey struct {
	Key string
}

type ctxTx struct {
	Tx *dbr.Tx
}

func (db *CtxDB) NewContextWithTx(ctx context.Context) context.Context {
	return context.WithValue(ctx, db.Key, &ctxTx{Tx: nil})
}

func (db *CtxDB) ctxTx(ctx context.Context) *ctxTx {
	ctxtx, ok := ctx.Value(db.Key).(*ctxTx)
	if !ok {
		panic(errors.New("No CtxDB Tx in context. Did you run NewContextWithTx"))
	}

	return ctxtx
}

func (db *CtxDB) Tx(ctx context.Context) *dbr.Tx {
	ctxtx := db.ctxTx(ctx)
	if ctxtx.Tx == nil {
		tx, err := db.NewSession(nil).BeginTx(ctx, nil)
		if err != nil {
			panic(errors.Wrap(err, "New CtxDB Tx failed."))
		}

		ctxtx.Tx = tx
	}

	return ctxtx.Tx
}

func (db *CtxDB) Commit(ctx context.Context) error {
	ctxtx := db.ctxTx(ctx)

	if ctxtx.Tx != nil {
		defer func() { ctxtx.Tx = nil }()
		return ctxtx.Tx.Commit()
	}

	return nil
}

func (db *CtxDB) CommitOrAbort(ctx context.Context) {
	if err := db.Commit(ctx); err != nil {
		panic(errors.Wrap(err, "CtxDB Tx Commit failed."))
	}
}

func (db *CtxDB) Rollback(ctx context.Context) error {
	ctxtx := db.ctxTx(ctx)

	if ctxtx.Tx != nil {
		defer func() { ctxtx.Tx = nil }()
		return ctxtx.Tx.Rollback()
	}

	return nil
}

func (db *CtxDB) RollbackUnlessCommitted(ctx context.Context) {
	ctxtx := db.ctxTx(ctx)

	if ctxtx.Tx != nil {
		defer func() { ctxtx.Tx = nil }()
		ctxtx.Tx.RollbackUnlessCommitted()
	}
}
